import numpy as np
from scipy import stats

def enumeration_method(p0=0.1, alpha=0.05, max_n=100):
    """
    枚举法求解最小样本量
    
    参数:
    p0: 标称次品率 (默认10%)
    alpha: 显著性水平 (默认0.05)
    max_n: 最大枚举样本量 (默认100)
    """
    
    # 计算z_(α/2)值
    z_alpha_2 = stats.norm.ppf(1 - alpha/2)
    print(f"z_(α/2) = {z_alpha_2:.4f}")
    print(f"标称次品率 p₀ = {p0}")
    print(f"显著性水平 α = {alpha}")
    print(f"条件: p̂ - {z_alpha_2:.4f} × SE > {p0}")
    print()
    
    print("=== 枚举法求解 ===")
    print(f"{'n':>4} {'k':>4} {'p̂':>8} {'SE':>10} {'条件值':>10} {'满足条件':>8}")
    print("-" * 60)
    
    # 1. 枚举 n = 1, 2, 3, ... 到 max_n
    for n in range(10, max_n + 1):
        # 2. 对每个 n, 尝试 k = 0 到 n 中的每一个整数
        for k in range(0, n + 1):
            # 计算 p̂ = k/n
            p_hat = k / n
            
            # 计算标准误 SE = √(p̂(1-p̂)/n)
            SE = np.sqrt(p_hat * (1 - p_hat) / n)
            
            # 计算条件值: p̂ - z_(α/2) × SE
            condition_value = p_hat - z_alpha_2 * SE
            
            # 判断是否满足条件: p̂ - z_(α/2) × SE > p₀
            satisfies_condition = condition_value > p0
            
            # 如果满足条件，输出结果
            if satisfies_condition:
                print(f"{n:4d} {k:4d} {p_hat:8.4f} {SE:10.6f} {condition_value:10.6f} {'是':>8}")
                print(f"\n找到满足条件的最小样本量: n = {n}")
                print(f"对应的k值: k = {k}")
                print(f"对应的p̂值: p̂ = {p_hat:.4f}")
                print(f"对应的SE值: SE = {SE:.6f}")
                print(f"条件值: {condition_value:.6f} > {p0}")
                return n, k, p_hat, SE
            else:
                # 可选：显示所有结果（注释掉以节省输出）
                # print(f"{n:4d} {k:4d} {p_hat:8.4f} {SE:10.6f} {condition_value:10.6f} {'否':>8}")
                pass
    
    print(f"\n在n ≤ {max_n}范围内未找到满足条件的样本量")
    return None

def analyze_specific_cases():
    """
    分析特定情况
    """
    print("\n=== 特定情况分析 ===")
    
    # 分析n=10的情况
    n = 10
    print(f"\n当n={n}时:")
    print(f"{'k':>4} {'p̂':>8} {'SE':>10} {'条件值':>10} {'满足条件':>8}")
    print("-" * 50)
    
    for k in range(0, n + 1):
        p_hat = k / n
        SE = np.sqrt(p_hat * (1 - p_hat) / n)
        condition_value = p_hat - 1.96 * SE
        satisfies = condition_value > 0.1
        
        print(f"{k:4d} {p_hat:8.4f} {SE:10.6f} {condition_value:10.6f} {'是' if satisfies else '否':>8}")

if __name__ == "__main__":
    # 使用枚举法求解
    result = enumeration_method(p0=0.1, alpha=0.05, max_n=1000)
    
    # 分析特定情况
    analyze_specific_cases()
